import requests
import numpy as np
from typing import List, Optional
import time
from rich.status import Status
from api_key import API_KEY

class EmbeddingAPI:
    def __init__(self):
        self.endpoint = "https://api.siliconflow.cn/v1/embeddings"
        self.headers = {
            "Authorization": f"Bearer {API_KEY}",
            "Content-Type": "application/json"
        }
        self.model = "BAAI/bge-m3"
        self.timeout = 30

    def get_embedding(self, text: str, status: Status = None) -> Optional[List[float]]:
        """获取单个文本的嵌入向量"""
        if status:
            status.update(f"正在处理: {text[:30]}...")
        
        for attempt in range(3):
            try:
                response = requests.post(
                    self.endpoint,
                    json={
                        "model": self.model,
                        "input": text,
                        "encoding_format": "float"
                    },
                    headers=self.headers,
                    timeout=self.timeout
                )
                response.raise_for_status()
                return response.json()['data'][0]['embedding']
            except Exception as e:
                if status:
                    status.update(f"[red]Attempt {attempt + 1} failed: {str(e)}[/]")
                time.sleep(min(2 ** attempt, 10))
        return None

    def batch_embed(self, texts: List[str], status: Status = None) -> List[Optional[List[float]]]:
        """批量获取嵌入"""
        if status:
            status.update(f"批量处理 {len(texts)} 条文本...")
        
        try:
            response = requests.post(
                self.endpoint,
                json={
                    "model": self.model,
                    "input": texts,
                    "encoding_format": "float"
                },
                headers=self.headers,
                timeout=self.timeout * 2
            )
            response.raise_for_status()
            return [item['embedding'] for item in response.json()['data']]
        except Exception as e:
            if status:
                status.update(f"[red]Batch failed: {str(e)}, 回退到单条模式[/]")
            return [self.get_embedding(text, status) for text in texts]
